#!/usr/bin/env python3
"""Various listeners for using with stomp.py connections."""

import logging
import sys
import threading
import time

from ambari_stomp.backward import monotonic
from ambari_stomp.constants import *
import ambari_stomp.exception as exception
import ambari_stomp.utils as utils


log = logging.getLogger("stomp.py")


class Publisher(object):
  """
  Simply a registry of listeners.
  """

  def set_listener(self, name, listener):
    """
    Set a named listener to use with this connection. See :py:class:`stomp.listener.ConnectionListener`

    :param str name: the name of the listener
    :param ConnectionListener listener: the listener object
    """
    pass

  def remove_listener(self, name):
    """
    Remove a listener.

    :param str name: the name of the listener to remove
    """
    pass

  def get_listener(self, name):
    """
    Return the named listener.

    :param str name: the listener to return

    :rtype: ConnectionListener
    """
    return None


class ConnectionListener(object):
  """
  This class should be used as a base class for objects registered
  using Connection.set_listener().
  """

  def on_connecting(self, host_and_port):
    """
    Called by the STOMP connection once a TCP/IP connection to the
    STOMP server has been established or re-established. Note that
    at this point, no connection has been established on the STOMP
    protocol level. For this, you need to invoke the "connect"
    method on the connection.

    :param (str,int) host_and_port: a tuple containing the host name and port number to which the connection
        has been established.
    """
    pass

  def on_connected(self, headers, body):
    """
    Called by the STOMP connection when a CONNECTED frame is
    received (after a connection has been established or
    re-established).

    :param dict headers: a dictionary containing all headers sent by the server as key/value pairs.
    :param body: the frame's payload. This is usually empty for CONNECTED frames.
    """
    pass

  def on_disconnected(self):
    """
    Called by the STOMP connection when a TCP/IP connection to the
    STOMP server has been lost.  No messages should be sent via
    the connection until it has been reestablished.
    """
    pass

  def on_heartbeat_timeout(self):
    """
    Called by the STOMP connection when a heartbeat message has not been
    received beyond the specified period.
    """
    pass

  def on_before_message(self, headers, body):
    """
    Called by the STOMP connection before a message is returned to the client app. Returns a tuple
    containing the headers and body (so that implementing listeners can pre-process the content).

    :param dict headers: the message headers
    :param body: the message body
    """
    return headers, body

  def on_message(self, headers, body):
    """
    Called by the STOMP connection when a MESSAGE frame is received.

    :param dict headers: a dictionary containing all headers sent by the server as key/value pairs.
    :param body: the frame's payload - the message body.
    """
    pass

  def on_receipt(self, headers, body):
    """
    Called by the STOMP connection when a RECEIPT frame is
    received, sent by the server if requested by the client using
    the 'receipt' header.

    :param dict headers: a dictionary containing all headers sent by the server as key/value pairs.
    :param body: the frame's payload. This is usually empty for RECEIPT frames.
    """
    pass

  def on_error(self, headers, body):
    """
    Called by the STOMP connection when an ERROR frame is received.

    :param dict headers: a dictionary containing all headers sent by the server as key/value pairs.
    :param body: the frame's payload - usually a detailed error description.
    """
    pass

  def on_send(self, frame):
    """
    Called by the STOMP connection when it is in the process of sending a message

    :param Frame frame: the frame to be sent
    """
    pass

  def on_heartbeat(self):
    """
    Called on receipt of a heartbeat.
    """
    pass


class HeartbeatListener(ConnectionListener):
  """
  Listener used to handle STOMP heartbeating.
  """

  def __init__(self, heartbeats):
    self.running = False
    self.heartbeats = heartbeats
    self.received_heartbeat = None
    self.heartbeat_thread = None
    self.next_outbound_heartbeat = None

  def on_connected(self, headers, body):
    """
    Once the connection is established, and 'heart-beat' is found in the headers, we calculate the real
    heartbeat numbers (based on what the server sent and what was specified by the client) - if the heartbeats
    are not 0, we start up the heartbeat loop accordingly.

    :param dict headers: headers in the connection message
    :param body: the message body
    """
    if "heart-beat" in headers:
      self.heartbeats = utils.calculate_heartbeats(
        headers["heart-beat"].replace(" ", "").split(","), self.heartbeats
      )
      if self.heartbeats != (0, 0):
        self.send_sleep = self.heartbeats[0] / 1000

        # receive gets an additional grace of 50%
        self.receive_sleep = (self.heartbeats[1] / 1000) * 1.5

        # Give grace of receiving the first heartbeat
        self.received_heartbeat = monotonic() + self.receive_sleep

        self.running = True
        if self.heartbeat_thread is None:
          self.heartbeat_thread = utils.default_create_thread(self.__heartbeat_loop)
          self.heartbeat_thread.name = "StompHeartbeat%s" % getattr(
            self.heartbeat_thread, "name", "Thread"
          )

  def on_disconnected(self):
    self.running = False

  def on_message(self, headers, body):
    """
    Reset the last received time whenever a message is received.

    :param dict headers: headers in the message
    :param body: the message content
    """
    # reset the heartbeat for any received message
    self.__update_heartbeat()

  def on_receipt(self, *_):
    """
    Reset the last received time whenever a receipt is received.
    """
    self.__update_heartbeat()

  def on_error(self, *_):
    """
    Reset the last received time whenever an error is received.
    """
    self.__update_heartbeat()

  def on_heartbeat(self):
    """
    Reset the last received time whenever a heartbeat message is received.
    """
    self.__update_heartbeat()

  def on_send(self, frame):
    """
    Add the heartbeat header to the frame when connecting, and bump
    next outbound heartbeat timestamp.

    :param Frame frame: the Frame object
    """
    if frame.cmd == CMD_CONNECT or frame.cmd == CMD_STOMP:
      if self.heartbeats != (0, 0):
        frame.headers[HDR_HEARTBEAT] = "%s,%s" % self.heartbeats
    if self.next_outbound_heartbeat is not None:
      self.next_outbound_heartbeat = monotonic() + self.send_sleep

  def __update_heartbeat(self):
    # Honour any grace that has been already included
    if self.received_heartbeat is None:
      return
    now = monotonic()
    if now > self.received_heartbeat:
      self.received_heartbeat = now

  def __heartbeat_loop(self):
    """
    Main loop for sending (and monitoring received) heartbeats.
    """
    now = monotonic()

    # Setup the initial due time for the outbound heartbeat
    if self.send_sleep != 0:
      self.next_outbound_heartbeat = now + self.send_sleep

    while self.running:
      now = monotonic()

      next_events = []
      if self.next_outbound_heartbeat is not None:
        next_events.append(self.next_outbound_heartbeat - now)
      if self.receive_sleep != 0:
        t = self.received_heartbeat + self.receive_sleep - now
        if t > 0:
          next_events.append(t)
      sleep_time = min(next_events)
      if sleep_time > 0:
        time.sleep(sleep_time)

      now = monotonic()

      if not self.transport.is_connected():
        time.sleep(self.send_sleep)
        continue

      if self.send_sleep != 0 and now > self.next_outbound_heartbeat:
        log.debug("Sending a heartbeat message at %s", now)
        try:
          self.transport.transmit(utils.Frame(None, {}, None))
        except exception.NotConnectedException:
          log.debug("Lost connection, unable to send heartbeat")
        except Exception:
          _, e, _ = sys.exc_info()
          log.debug("Unable to send heartbeat, due to: %s", e)

      if self.receive_sleep != 0:
        diff_receive = now - self.received_heartbeat

        if diff_receive > self.receive_sleep:
          # heartbeat timeout
          log.warning(
            "Heartbeat timeout: diff_receive=%s, time=%s, lastrec=%s",
            diff_receive,
            now,
            self.received_heartbeat,
          )
          self.transport.set_connected(False)
          self.transport.disconnect_socket()
          self.transport.stop()
          for listener in self.transport.listeners.values():
            listener.on_heartbeat_timeout()
    self.heartbeat_thread = None


class WaitingListener(ConnectionListener):
  """
  A listener which waits for a specific receipt to arrive.
  """

  def __init__(self, receipt):
    """
    :param str receipt:
    """
    self.condition = threading.Condition()
    self.receipt = receipt
    self.received = False

  def on_receipt(self, headers, body):
    """
    If the receipt id can be found in the headers, then notify the waiting thread.

    :param dict headers: headers in the message
    :param body: the message content
    """
    if "receipt-id" in headers and headers["receipt-id"] == self.receipt:
      with self.condition:
        self.received = True
        self.condition.notify()

  def wait_on_receipt(self):
    """
    Wait until we receive a message receipt.
    """
    with self.condition:
      while not self.received:
        self.condition.wait()
    self.received = False


class StatsListener(ConnectionListener):
  """
  A connection listener for recording statistics on messages sent and received.
  """

  def __init__(self):
    # The number of errors received
    self.errors = 0
    # The number of connections established
    self.connections = 0
    # The number of disconnections
    self.disconnects = 0
    # The number of messages received
    self.messages = 0
    # The number of messages sent
    self.messages_sent = 0
    # The number of heartbeat timeouts
    self.heartbeat_timeouts = 0
    # The number of heartbeats
    self.heartbeat_count = 0

  def on_disconnected(self):
    """
    Increment the disconnect count. See :py:meth:`ConnectionListener.on_disconnected`
    """
    self.disconnects += 1
    log.info("disconnected (x %s)", self.disconnects)

  def on_error(self, headers, body):
    """
    Increment the error count. See :py:meth:`ConnectionListener.on_error`

    :param dict headers: headers in the message
    :param body: the message content
    """
    log.info("received an error %s [%s]", body, headers)
    self.errors += 1

  def on_connecting(self, host_and_port):
    """
    Increment the connection count. See :py:meth:`ConnectionListener.on_connecting`

    :param (str,int) host_and_port: the host and port as a tuple
    """
    log.info(
      "connecting %s %s (x %s)", host_and_port[0], host_and_port[1], self.connections
    )
    self.connections += 1

  def on_message(self, headers, body):
    """
    Increment the message received count. See :py:meth:`ConnectionListener.on_message`

    :param dict headers: headers in the message
    :param body: the message content
    """
    self.messages += 1

  def on_send(self, frame):
    """
    Increment the send count. See :py:meth:`ConnectionListener.on_send`

    :param Frame frame:
    """
    self.messages_sent += 1

  def on_heartbeat_timeout(self):
    """
    Increment the heartbeat timeout. See :py:meth:`ConnectionListener.on_heartbeat_timeout`
    """
    log.debug("received heartbeat timeout")
    self.heartbeat_timeouts += 1

  def on_heartbeat(self):
    """
    Increment the heartbeat count. See :py:meth:`ConnectionListener.on_heartbeat`
    """
    self.heartbeat_count += 1

  def __str__(self):
    """
    Return a string containing the current statistics (messages sent and received,
    errors, etc)
    """
    return f"""Connections: {self.connections}
Messages sent: {self.messages_sent}
Messages received: {self.messages}
Heartbeats received: {self.heartbeat_count}
Errors: {self.errors}"""


class PrintingListener(ConnectionListener):
  def on_connecting(self, host_and_port):
    """
    :param (str,int) host_and_port:
    """
    print("on_connecting %s %s" % host_and_port)

  def on_connected(self, headers, body):
    """
    :param dict headers:
    :param body:
    """
    print(f"on_connected {headers} {body}")

  def on_disconnected(self):
    print("on_disconnected")

  def on_heartbeat_timeout(self):
    print("on_heartbeat_timeout")

  def on_before_message(self, headers, body):
    """
    :param dict headers:
    :param body:
    """
    print(f"on_before_message {headers} {body}")
    return headers, body

  def on_message(self, headers, body):
    """
    :param dict headers:
    :param body:
    """
    print(f"on_message {headers} {body}")

  def on_receipt(self, headers, body):
    """
    :param dict headers:
    :param body:
    """
    print(f"on_receipt {headers} {body}")

  def on_error(self, headers, body):
    """
    :param dict headers:
    :param body:
    """
    print(f"on_error {headers} {body}")

  def on_send(self, frame):
    """
    :param Frame frame:
    """
    print(f"on_send {frame.cmd} {frame.headers} {frame.body}")

  def on_heartbeat(self):
    print("on_heartbeat")


class TestListener(StatsListener, WaitingListener):
  """
  Implementation of StatsListener and WaitingListener. Useful for testing.
  """

  def __init__(self, receipt=None):
    """
    :param str receipt:
    """
    StatsListener.__init__(self)
    WaitingListener.__init__(self, receipt)
    self.message_list = []
    self.message_condition = threading.Condition()
    self.message_received = False
    self.heartbeat_condition = threading.Condition()
    self.heartbeat_received = False

  def on_message(self, headers, message):
    """
    :param dict headers:
    :param message:
    """
    StatsListener.on_message(self, headers, message)
    self.message_list.append((headers, message))
    with self.message_condition:
      self.message_received = True
      self.message_condition.notify()

  def wait_for_message(self):
    with self.message_condition:
      while not self.message_received:
        self.message_condition.wait()
    self.message_received = False

  def get_latest_message(self):
    return self.message_list[-1]

  def on_heartbeat(self):
    StatsListener.on_heartbeat(self)
    with self.heartbeat_condition:
      self.heartbeat_received = True
      self.heartbeat_condition.notify()

  def wait_for_heartbeat(self):
    with self.heartbeat_condition:
      while not self.heartbeat_received:
        self.heartbeat_condition.wait()
    self.heartbeat_received = False
